"""
Examples of different loss function configurations for the continual learning framework.

This file shows how to configure different types of loss functions in your training arguments.
"""

# Example 1: Cross Entropy Loss (Simple classification)
cross_entropy_config = {
    "loss_type": "cross_entropy",
    # No additional parameters needed
}

# Example 1b: Weighted Cross Entropy Loss (with scaling factor)
weighted_cross_entropy_config = {
    "loss_type": "cross_entropy",
    "scale": 2.0,  # Scale cross entropy loss by factor of 2
}

# Example 2: CosFace Loss (Angular margin loss)
cosface_config = {
    "loss_type": "cosface",
    "scale": 20.0,  # Scaling factor
    "margin": 0.0,  # Margin parameter
    "eps": 1e-7,  # Small epsilon for numerical stability
}

# Example 3: ArcFace Loss (Additive angular margin)
arcface_config = {
    "loss_type": "arcface",
    "scale": 64.0,  # Scaling factor
    "margin": 0.5,  # Angular margin
    "eps": 1e-7,
}

# Example 4: SphereFace Loss (Multiplicative angular margin)
sphereface_config = {
    "loss_type": "sphereface",
    "scale": 64.0,
    "margin": 1.35,
    "eps": 1e-7,
}

# Example 5: Focal Loss (For imbalanced datasets)
focal_loss_config = {
    "loss_type": "focal",
    "focal_alpha": 1.0,  # Weighting factor
    "focal_gamma": 2.0,  # Focusing parameter
}

# Example 6: Label Smoothing Loss (Regularization technique)
label_smoothing_config = {
    "loss_type": "label_smoothing",
    "label_smoothing": 0.1,  # Smoothing factor (0.0 = no smoothing, 1.0 = uniform)
}

# Example 7: Knowledge Distillation Loss (For continual learning)
distillation_config = {
    "loss_type": "distillation",
    "distill_temperature": 4.0,  # Temperature for softmax
    "alpha": 0.5,  # Balance between hard and soft targets
}


def get_loss_config(loss_name):
    """
    Get a specific loss configuration by name.

    Args:
        loss_name (str): Name of the loss configuration

    Returns:
        dict: Loss configuration parameters
    """
    configs = {
        "cross_entropy": cross_entropy_config,
        "weighted_cross_entropy": weighted_cross_entropy_config,
        "cosface": cosface_config,
        "arcface": arcface_config,
        "sphereface": sphereface_config,
        "focal": focal_loss_config,
        "label_smoothing": label_smoothing_config,
        "distillation": distillation_config,
    }

    if loss_name not in configs:
        raise ValueError(f"Unknown loss config: {loss_name}")

    return configs[loss_name]


def print_available_losses():
    """Print all available loss configurations"""
    from utils.loss import LossFactory

    print("Available loss functions:")
    for loss_type in LossFactory.get_available_losses():
        print(f"  - {loss_type}")

    print("\nExample configurations:")
    print("  - cross_entropy: Simple cross-entropy loss")
    print("  - cosface: Cosine face loss with margin")
    print("  - arcface: Additive angular margin loss")
    print("  - sphereface: Multiplicative angular margin loss")
    print("  - focal: Focal loss for imbalanced classes")
    print("  - label_smoothing: Label smoothing regularization")
    print("  - distillation: Knowledge distillation loss")


if __name__ == "__main__":
    print_available_losses()

    # Example usage
    print("\nExample usage in main training script:")
    print("args.update(get_loss_config('focal'))")
    print("# This will use focal loss with alpha=1.0, gamma=2.0")
